# -*- coding: utf-8 -*-

# cite: https://gitee.com/ascend/MindSpeed/blob/master/mindspeed/patch_utils.py

import sys
import types
import importlib
from typing import List, Dict, Union, Callable


def dontexist_func_wrapper(func_name: str):
    def dontexist_func(*args, **kwargs):
        raise RuntimeError('function {} no exist'.format(func_name))
    return dontexist_func


class Patch:
    def __init__(self,
                 func_full_name: str,
                 new_func: Callable = None,
                 create_dummy: bool = False):
        func_path = func_full_name.rsplit('.', 1)
        if len(func_path) == 1:
            self.module_path, self.func_full_name = func_full_name, None
        else:
            self.module_path, self.func_full_name = func_path
        self.module = None
        self.ori_func = None
        self.patch_func = None
        self.wrappers = []
        if new_func is None:
            new_func = dontexist_func_wrapper(func_full_name)
        self.set_patch_func(new_func)
        self.is_applied = False
        self.create_dummy = create_dummy
        
    def set_patch_func(self, new_func: Callable, force_patch=False):
        if hasattr(new_func, '__name__') and new_func.__name__.endswith(('wrapper', 'decorator')):
            if new_func not in self.wrappers:
                self.wrappers.append(new_func)
        else:
            if self.patch_func and not force_patch:
                raise RuntimeError('the patch of {} exist !'.format(self.func_full_name))
            self.patch_func = new_func
        self.is_applied = False

    @property
    def ori_func_id(self):
        return id(self.ori_func)

    def remove_wrappers(self, wrapper_names: Union[str, List[str]] = None):
        if wrapper_names is None:
            self.wrappers.clear()
            return
        if isinstance(wrapper_names, str):
            wrapper_names = [wrapper_names]
        for name in wrapper_names:
            i = 0
            while i < len(self.wrappers):
                if self.wrappers[i].__name__ == name:
                    self.wrappers.pop(i)
                else:
                    i += 1
                    
    @staticmethod
    def parse_path(module_path: str, func_name: str, create_dummy: bool):
        from importlib.machinery import ModuleSpec
        modules = module_path.split('.')
        for i in range(1, len(modules) + 1):
            parent = '.'.join(modules[:i - 1])
            path = '.'.join(modules[:i])
            try:
                importlib.import_module(path)
            except ModuleNotFoundError as e:
                if not parent or not hasattr(importlib.import_module(parent), modules[i - 1]):
                    if not create_dummy:
                        raise ModuleNotFoundError(e) from e
                    sys.modules[path] = types.ModuleType(path)
                    sys.modules[path].__file__ = 'dramkit.dummy_module.py'
                    sys.modules[path].__spec__ = ModuleSpec(path, None)
                    if parent:
                        setattr(importlib.import_module(parent), modules[i - 1], sys.modules[path])
                else:
                    module = getattr(importlib.import_module(parent), modules[i - 1])
                    if hasattr(module, func_name):
                        return module, getattr(module, func_name)
                    elif create_dummy:
                        return module, dontexist_func_wrapper(func_name)
                    else:
                        raise RuntimeError('no exist {} of {}'.format(func_name, module))
        if func_name is not None and not hasattr(sys.modules[module_path], func_name):
            setattr(sys.modules[module_path], func_name, None)
        return sys.modules[module_path], getattr(sys.modules[module_path], func_name) if func_name is not None else None


    def apply(self):
        if self.is_applied:
            return

        self.module, self.ori_func = Patch.parse_path(self.module_path, self.func_full_name, self.create_dummy)

        final_func = self.ori_func
        if self.patch_func is not None:
            final_func = self.patch_func

        for wrapper in self.wrappers:
            final_func = wrapper(final_func)

        if self.func_full_name is not None:
            setattr(self.module, self.func_full_name, final_func)
        for key, value in sys.modules.copy().items():
            if self.func_full_name is not None \
                    and hasattr(value, self.func_full_name) \
                    and id(getattr(value, self.func_full_name)) == self.ori_func_id:
                setattr(value, self.func_full_name, final_func)

        self.is_applied = True


class PatchesManager:
    patches_info: Dict[str, Patch] = {}

    @staticmethod
    def register_patch(func_full_name: str,
                       new_func: Callable = None,
                       create_dummy: bool = False,
                       force_patch: bool = False):
        if func_full_name not in PatchesManager.patches_info:
            PatchesManager.patches_info[func_full_name] = \
                Patch(func_full_name, new_func, create_dummy)
        else:
            PatchesManager.patches_info.get(func_full_name).set_patch_func(
                new_func, force_patch)

    @staticmethod
    def remove_wrappers(func_full_name: str,
                        wrapper_names: Union[str, List[str]],
                        remove_check: bool = True):
        if func_full_name not in PatchesManager.patches_info:
            raise ValueError('The function <{}> not exist.'.format(func_full_name))
        patch = PatchesManager.patches_info.get(func_full_name)
        wrappers_len = len(patch.wrappers)
        patch.remove_wrappers(wrapper_names)
        if remove_check and wrappers_len == len(patch.wrappers):
            raise RuntimeError('Remove wrappers has not remove anything.')

    @staticmethod
    def apply_patches():
        for patch in PatchesManager.patches_info.values():
            patch.apply()
